import common.ConstVars;
import image.ImageCapturer;
import image.ImageCapturerImpl;
import image.ImageFixer;

import javax.swing.*;
import java.awt.*;
import java.awt.event.ActionEvent;
import java.awt.event.ActionListener;
import java.awt.event.KeyEvent;
import java.awt.event.KeyListener;
import java.io.*;
import java.awt.image.BufferedImage;

public class TrainingImageCapturer extends JFrame {

    protected Button alex, michael, kevin, aaron, other, save;
    protected JLabel picture;
    private ImageCapturer ic;
    private String curName;
    private static String tempOutput = "temp.jpg";
    private static String tempIcon = "ticon.jpg";
    private static String dir;

    // Need to pass robot into constructor
    public TrainingImageCapturer(String dir)
    {
        super("Training Image Capturer");

        this.dir = dir;

        setSize(650, 600);
        this.setFocusable(true);

        KeyListener kl = new MyKeyListener();

        alex = new Button(ConstVars.ALEX);
        michael = new Button(ConstVars.MICHAEL);
        kevin = new Button(ConstVars.KEVIN);
        aaron = new Button(ConstVars.AARON);
        other = new Button(ConstVars.OTHER);
        save = new Button("Save");

        this.addKeyListener(kl);
        alex.addKeyListener(kl);
        michael.addKeyListener(kl);
        kevin.addKeyListener(kl);
        aaron.addKeyListener(kl);
        other.addKeyListener(kl);
        save.addKeyListener(kl);

        ButtonListener b = new ButtonListener();

        alex.addActionListener(b);
        michael.addActionListener(b);
        kevin.addActionListener(b);
        aaron.addActionListener(b);
        other.addActionListener(b);
        save.addActionListener(b);

        setLayout(new BorderLayout());
        Panel p2 = new Panel();
        add(p2, BorderLayout.NORTH);

        p2.setLayout(new BorderLayout());
        p2.add(alex, BorderLayout.NORTH);
        p2.add(michael, BorderLayout.SOUTH);
        p2.add(kevin, BorderLayout.WEST);
        p2.add(aaron, BorderLayout.EAST);
        p2.add(other, BorderLayout.CENTER);

        picture = new JLabel("PRESS A BUTTON ABOVE");
        picture.addKeyListener(kl);
        picture.setFocusable(true);
        add(picture, BorderLayout.CENTER);

        add(save, BorderLayout.SOUTH);

        ic = new ImageCapturerImpl();

        this.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
        setVisible(true);
    }

    public static void copyFile(File in, File out) throws Exception {
        FileInputStream fis  = new FileInputStream(in);
        FileOutputStream fos = new FileOutputStream(out);
        try {
            byte[] buf = new byte[1024];
            int i = 0;
            while ((i = fis.read(buf)) != -1) {
                fos.write(buf, 0, i);
            }
        }
        catch (Exception e) {
            throw e;
        }
        finally {
            if (fis != null) fis.close();
            if (fos != null) fos.close();
        }
    }

    protected void saveImage()
    {
        if(dir.charAt(dir.length()-1) != File.separatorChar)
            dir += File.separatorChar;
        File fDir = new File(dir + curName + File.separatorChar);
        int num = fDir.list().length;

        String outFile = String.format("%s%s%c%04d.jpg", dir, curName, File.separatorChar, num);

        try {
            copyFile(new File(tempOutput), new File(outFile));
        }
        catch(Exception e) {
            e.printStackTrace();
        }
    }

    private class MyKeyListener implements KeyListener {

        public void keyTyped(KeyEvent e) {
            if(e.getKeyCode() == KeyEvent.VK_ENTER) {
                saveImage();
            }
            else if(e.getKeyCode() == KeyEvent.VK_W) {
                takePictureFor(ConstVars.ALEX);
            }
            else if(e.getKeyCode() == KeyEvent.VK_A) {
                takePictureFor(ConstVars.KEVIN);
            }
            else if(e.getKeyCode() == KeyEvent.VK_S) {
                takePictureFor(ConstVars.MICHAEL);
            }
            else if(e.getKeyCode() == KeyEvent.VK_D) {
                takePictureFor(ConstVars.AARON);
            }
            else if(e.getKeyCode() == KeyEvent.VK_SPACE) {
                takePictureFor(ConstVars.OTHER);
            }
        }

        public void keyPressed(KeyEvent e) {
            //To change body of implemented methods use File | Settings | File Templates.
        }

        public void keyReleased(KeyEvent e) {
            //To change body of implemented methods use File | Settings | File Templates.
        }
    }

    private void updateImage(String fileName)
    {
        ImageIcon i = (ImageIcon)picture.getIcon();
        if(i != null)
            i.getImage().flush();
        picture.setIcon(new ImageIcon(fileName));
        this.repaint();
    }

    private class ButtonListener implements ActionListener
    {
        public void actionPerformed(ActionEvent evt)
        {
            String arg = evt.getActionCommand();
            if(arg.equals("Save"))
            {
                saveImage();
            }
            else
            {
                takePictureFor(arg);
            }
        }
    }

    private void takePicture()
    {
        ic.captureImage(tempOutput);
        BufferedImage bimage = null;
        try {
            bimage = ImageFixer.getImageFromFile(tempOutput);
            ImageFixer.drawFeaturePoints(bimage);
            ImageFixer.save(bimage, tempIcon);
            updateImage(tempIcon);
            picture.setText("");
            picture.updateUI();
        }
        catch(Exception e) {
            System.err.println("image read failed in takePicture; frame skipped");
//            e.printStackTrace();
        }
    }

    private void takePictureFor(String arg)
    {
        curName = arg;
    }

    public static void main(String args[])
    {
        if(args.length == 0)
        {
            System.err.println("Dir name????");
            System.exit(1);
        }
        TrainingImageCapturer tic = new TrainingImageCapturer(args[0]);
        tic.setVisible(true);
        int c = 0;
        while(true) {
            tic.takePicture();
        }
    }
}
